#include "SaveAsDataset.h"

#include <unistd.h>
#include <fcntl.h>

SaveAsDataset::SaveAsDataset(std::string save_path)
    : save_path(std::move(save_path))
{

}

void SaveAsDataset::draw_armors(cv::Mat &img, SaveAsDataset::Armors &armors, int delay)
{
    for (auto &armor: armors)
    {
        cv::rectangle(img, armor.rect(), [&]()
        {
            switch (armor.armor_color)
            {
                case Armor::Color::RED: return cv::Scalar(0, 0, 255);
                case Armor::Color::BLUE: return cv::Scalar(255, 200, 0);
                default: return cv::Scalar(255, 255, 255);
            }
        }(), 1);
    }
    if (delay >= 0)
    {
        cv::imshow("image", img);
        cv::waitKey(delay);
    }
}

int SaveAsDataset::correct_label(const cv::Mat &img, int &label, int exist_label, int i)
{
    if (exist_label != -1)
        label = exist_label;
    int key = 1;
    cv::Mat large_img;
    cv::resize(img, large_img, cv::Size(), 10, 10);
    while (key != 0)
    {
        cv::circle(large_img, cv::Point(22, 37), 20, cv::Scalar(0, 0, 0), cv::FILLED);
        cv::putText(large_img, std::to_string(label), cv::Point(10, 50), cv::FONT_HERSHEY_DUPLEX, 1.2, cv::Scalar(255, 255, 255));
        cv::imshow("armor", large_img);
        key = cv::waitKey(0);
        switch (key)
        {
            case 'a':--i;
                key = 0;
                break;
            case 'd':++i;
                key = 0;
                break;
            case 177:
            case '1':label = 1;
                break;
            case 178:
            case '2':label = 2;
                break;
            case 179:
            case '3':label = 3;
                break;
            case 180:
            case '4':label = 4;
                break;
            case 181:
            case '5':label = 5;
                break;
            case 182:
            case '6':label = 6;
                break;
            case 183:
            case '7':label = 7;
                break;
            case 184:
            case '8':label = 8;
                break;
            case 's':
            case 176:
            case '0':label = 0;
                break;
            default:std::cerr << '\r' << "Please press an valid key!  [" << key << "]" << std::endl;
                break;
        }
    }
    return i;
}

int
SaveAsDataset::label_and_save(int img_cnt, bool back, const std::initializer_list<Armors> &armor_list, bool save_error_only)
{
    Armors armors;
    collect_armors(armor_list, armors);
    int nxt_i = back ? -1 : 1; // 考虑如果armors为空，那么下一个img的索引由nxt_i决定。
    for (int i = back ? (armors.size() - 1) : 0; i < armors.size(); i = nxt_i)
    {
        std::cerr << '\r' << "look at armor: img[" << img_cnt << "]::armor[" << i << "]";
        auto exist_label = find_exist_label(img_cnt, i);
        auto old_label = armors[i].classifier_num;
        nxt_i = correct_label(armors[i].number_img, armors[i].classifier_num, exist_label, i);
        // 若只保存分类错误，那么需要标签发生变化才保存，除非当前照片已经存在且不等于修改结果
        if (!save_error_only || old_label != armors[i].classifier_num
            || (exist_label != -1 && exist_label != armors[i].classifier_num))
            write_img(armors[i].number_img, img_cnt, i, armors[i].classifier_num, exist_label);
    }
    return nxt_i < 0 ? --img_cnt : ++img_cnt;
}

void
SaveAsDataset::collect_armors(const std::initializer_list<Armors> &armor_list, SaveAsDataset::Armors &armors)
{
    size_t sum_armor = 0;
    std::for_each_n(begin(armor_list), armor_list.size(), [&](const Armors &p)
    {
        sum_armor += p.size();
    });
    armors.clear();
    armors.reserve(sum_armor);
    for (const auto &armors_arr: armor_list)
    {
        for (const auto &armor: armors_arr)
            armors.emplace_back(armor);
    }
}

void
SaveAsDataset::write_img(const cv::Mat &img, int img_cnt, int armor_cnt, int label, int exist_label)
{
    auto filename =
        save_path + "autogen_" + std::to_string(img_cnt) + "_" + std::to_string(armor_cnt)
        + "_" + std::to_string(label) + ".bmp";
    if (exist_label == -1)
    {
        cv::imwrite(filename, img);
        std::cerr << '\r' << "save file: " << filename << std::endl;
    }
    else if (exist_label != label)
    {
        auto oldfile =
            save_path + "autogen_" + std::to_string(img_cnt) + "_" + std::to_string(armor_cnt)
            + "_" + std::to_string(exist_label) + ".bmp";
        if (0 != rename(oldfile.c_str(), filename.c_str()))
            throw std::runtime_error("rename failed!");
        std::cerr << '\r' << "relabel file: " << filename << std::endl;
    }
}

int SaveAsDataset::find_exist_label(int img_cnt, int armor_cnt)
{
    auto file_name =
        save_path + "autogen_" + std::to_string(img_cnt) + "_" + std::to_string(armor_cnt) + "_*"
        + ".bmp";
    auto num_pos = file_name.rfind('*');
    int old_label;
    for (old_label = 8; old_label >= 0; --old_label)
    {
        file_name[num_pos] = old_label + '0';
        if (access(file_name.c_str(), F_OK) == 0)
        {
            break;
        }
    }
    return old_label;
}
